// author: Shelman
// description: 封装的使用posix接口实现的UDP工具类

#include "UdpUtil.h"

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#pragma comment(lib, "ws2_32.lib")
#else
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>
#endif

#include "Common/CommLogDef.h"



// 填写地址信息
static bool packageAddrInfo(const NetworkConnInfo &stConnInfo, struct sockaddr_in &stAddrInfo)
{
    memset(&stAddrInfo, 0, sizeof(stAddrInfo));
    stAddrInfo.sin_family = AF_INET;
    stAddrInfo.sin_port = htons(stConnInfo.ui16Port);

    bool bRet = true;
    if (std::string("0.0.0.0") == stConnInfo.strIP || stConnInfo.strIP.empty()) {
        stAddrInfo.sin_addr.s_addr = INADDR_ANY;
    } else {
#ifdef _WIN32
        stAddrInfo.sin_addr.s_addr = inet_addr(stConnInfo.strIP.c_str());
        if (stAddrInfo.sin_addr.s_addr == INADDR_NONE) {
            COMMONLOG_ERROR("Invalid IP address: %s.", stConnInfo.strIP.data());
            return false;
        }
#else
        if (inet_aton(stConnInfo.strIP.c_str(), &stAddrInfo.sin_addr) == 0) {
            COMMONLOG_ERROR("Invalid IP address: %s.", stConnInfo.strIP.data());
            return false;
        }
#endif
    }

    return bRet;
}

// 解析地址信息
static bool parseAddrInfo(const struct sockaddr_in &stAddrInfo, NetworkConnInfo &stConnInfo)
{
    stConnInfo.ui16Port = ntohs(stAddrInfo.sin_port);
#ifdef _WIN32
    char* pClientAddr = inet_ntoa(stAddrInfo.sin_addr);
    if (pClientAddr) {
        stConnInfo.strIP = std::string(pClientAddr);
    } else {
        return false;
    }
#else
    char szClientIP[64] = {0};
    if (!inet_ntop(AF_INET, &stAddrInfo.sin_addr, szClientIP, sizeof(szClientIP))) {
        return false;
    }
    stConnInfo.strIP = std::string(szClientIP);
#endif
    return true;
}


/*****************************************************
 * 功能：构造函数
 * ***************************************************/
UdpUtil::UdpUtil()
{
#ifdef _WIN32
    // Windows Socket初始化
    WSADATA wsaData;
    WSAStartup(MAKEWORD(2, 2), &wsaData);
#endif
}

/*****************************************************
 * 功能：析构函数
 * ***************************************************/
UdpUtil::~UdpUtil()
{
#ifdef _WIN32
    WSACleanup();
#endif
}

/*****************************************************
 * 功能：设置接收数据的回调，
 * 输入参数：
 *      pfnCallback：数据回调函数
 * ***************************************************/
void UdpUtil::setDataRecvCallback(const UdpDataReceivedCallback &pfnCallback)
{
    std::lock_guard<std::mutex> lock(m_mt4PfnCallback);
    m_pfn4RecvDataCb = pfnCallback;
    return;
}

/*****************************************************
 * 功能：开启UDP端口
 * 输入参数：
 *      strLocalInfo：本地连接信息
 * 返回值：
 *      bool：开启结果，true -- 成功，false -- 失败
 * ***************************************************/
bool UdpUtil::start(const NetworkConnInfo &strLocalInfo)
{
    if (m_bRunning) {
        COMMONLOG_ERROR("Local udp (%s:%u) has been started, please close it first."
                        , m_stLocalConnInfo.strIP.data(), m_stLocalConnInfo.ui16Port);
        return false;
    } else {
        // 创建UDP套接字
        m_iLocalSockFd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
        if (0 > m_iLocalSockFd) {
            COMMONLOG_ERROR("%s", "Create socket failed.");
            return false;
        }

        m_stLocalConnInfo = strLocalInfo;

        // 设置本地地址
        struct sockaddr_in stLocalAddr;
        if (!packageAddrInfo(m_stLocalConnInfo, stLocalAddr)) {
            closeSocket();
            return false;
        }

        // 绑定套接字
        if (0 > bind(m_iLocalSockFd, (struct sockaddr*)&stLocalAddr, sizeof(stLocalAddr))) {
            COMMONLOG_ERROR("Bind udp socket failed, local (%s:%u)."
                            , m_stLocalConnInfo.strIP.data(), m_stLocalConnInfo.ui16Port);
            closeSocket();
            return false;
        }

        // 获取实际绑定的端口号
#ifdef __linux__
        socklen_t sAddrLen = sizeof(stLocalAddr);
        if (0 > getsockname(m_iLocalSockFd, (struct sockaddr*)&stLocalAddr, &sAddrLen)) {
            COMMONLOG_WARN("Getsockname failed, local (%s:%u)."
                           , m_stLocalConnInfo.strIP.data(), m_stLocalConnInfo.ui16Port);
        }
#else
        int iAddrLen = sizeof(stLocalAddr);
        if (0 > getsockname(m_iLocalSockFd, (struct sockaddr*)&stLocalAddr, &iAddrLen)) {
            COMMONLOG_WARN("Getsockname failed, local (%s:%u)."
                           , m_stLocalConnInfo.strIP.data(), m_stLocalConnInfo.ui16Port);
        }
#endif

        // 设置receive超时时间
        struct timeval stRecvTimeout = {5, 0};   //{秒，微妙}
        setsockopt(m_iLocalSockFd, SOL_SOCKET, SO_RCVTIMEO, (char*)(&stRecvTimeout), sizeof(stRecvTimeout));

        m_bRunning = true;

        // 开启数据接收线程
        m_th4RecvThread = std::thread(&UdpUtil::onRecvDataThread, this);
        COMMONLOG_INFO("Local udp start on (%s:%u).", m_stLocalConnInfo.strIP.data(), m_stLocalConnInfo.ui16Port);
    }

    return true;
}

/*****************************************************
 * 功能：停止UDP
 * ***************************************************/
void UdpUtil::stop()
{
    if (m_bRunning) {
        closeSocket();

        m_bRunning = false;
        if (m_th4RecvThread.joinable()) {
            m_th4RecvThread.join();
        }
    }

    return;
}

/*****************************************************
 * 功能：向远端发送数据
 * 输入参数：
 *      stRemoteInfo：远端信息
 *      strData：待发送的数据
 * 返回值：
 *      bool：发送结果，true -- 发送成功，false -- 发送失败
 * ***************************************************/
bool UdpUtil::sendData(const NetworkConnInfo &stRemoteInfo, const std::string &strData)
{
    std::lock_guard<std::mutex> lock(m_mt4SendData);
    if (!m_bRunning || stRemoteInfo.strIP.empty() || 0 >= stRemoteInfo.ui16Port) {
        return false;
    }

    struct sockaddr_in stAddrInfo;
    if (!packageAddrInfo(stRemoteInfo, stAddrInfo)) {
        return false;
    }

    bool bRet = true;
    size_t sTotalSent = 0;
    size_t sDataLen = strData.length();

    while (sTotalSent < sDataLen) {
        int iSent = sendto(m_iLocalSockFd, strData.c_str() + sTotalSent, sDataLen - sTotalSent
                           , 0, (struct sockaddr*)&stAddrInfo, sizeof(stAddrInfo));
        if (0 <= iSent) {
            sTotalSent += iSent;
        } else {
#ifdef _WIN32
            int error = WSAGetLastError();
            if (error == WSAEWOULDBLOCK) {
                // 非阻塞模式下缓冲区满，稍等继续发送
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            } else {
                COMMONLOG_ERROR("Failed to send data: %d.", error);
                bRet = false;
                break;
            }
#else
            if (errno == EWOULDBLOCK || errno == EAGAIN) {
                // 非阻塞模式下缓冲区满，稍等继续发送
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
                continue;
            } else {
                COMMONLOG_ERROR("Failed to send data: %s.", strerror(errno));
                bRet = false;
                break;
            }
#endif
        }
    }
    return bRet;
}

/*****************************************************
 * 功能：处理接收数据的线程
 * ***************************************************/
void UdpUtil::onRecvDataThread()
{
    COMMONLOG_INFO("%s", "Enter in udp recv data thread.");

    // 设置select()参数
    struct timeval timeout{};
    timeout.tv_sec = 3; // 设置3秒超时
    timeout.tv_usec = 0;

    char szBuffer[RECVDATA_DEFAULT_SIZE] = {0};
    fd_set readfds; // 数据读取fds

    while (m_bRunning) {
        FD_ZERO(&readfds);
        FD_SET(m_iLocalSockFd, &readfds);

        int result = select(m_iLocalSockFd + 1, &readfds, nullptr, nullptr, &timeout);
        if (result < 0) {
            COMMONLOG_ERROR("%s", "Select socket read event error!");
            break;
        } else if (0 == result) {
            // 超时，数据可读，继续循环
            std::this_thread::sleep_for(std::chrono::milliseconds(5)); // 让出CPU时间
            continue;
        }

        // 处理读操作，读取数据
        if (FD_ISSET(m_iLocalSockFd, &readfds)) {
            memset(szBuffer, 0, sizeof(szBuffer));

            struct sockaddr_in stRemoteAddr;

#ifdef __linux__
            socklen_t iAddrLen = sizeof(stRemoteAddr);
            int iRecvLen = recvfrom(m_iLocalSockFd, szBuffer, sizeof(szBuffer), 0
                                    , (struct sockaddr*)&stRemoteAddr, &iAddrLen);
#else
            int iAddrLen = sizeof(stRemoteAddr);
            int iRecvLen = recvfrom(m_iLocalSockFd, szBuffer, sizeof(szBuffer), 0
                                    , (struct sockaddr*)&stRemoteAddr, &iAddrLen);
#endif
            if (0 < iRecvLen) {
                NetworkConnInfo stConnInfo;
                if (!parseAddrInfo(stRemoteAddr, stConnInfo)) {
                    COMMONLOG_WARN("%s", "Parse remote addrinfo failed.");
                    std::this_thread::sleep_for(std::chrono::milliseconds(5)); // 让出CPU时间
                    continue;
                }

                std::string strData(szBuffer, static_cast<size_t>(iRecvLen));
                m_mt4PfnCallback.lock();
                if (m_pfn4RecvDataCb) {
                    m_pfn4RecvDataCb(strData, stConnInfo);
                }
                m_mt4PfnCallback.unlock();
            }
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(10)); // 让出CPU时间
    }

    COMMONLOG_INFO("%s", "Exit to udp recv data thread.");
    return;
}

/*****************************************************
 * 功能：关闭socket
 * ***************************************************/
void UdpUtil::closeSocket()
{
    if (0 <= m_iLocalSockFd) {
#ifdef _WIN32
        closesocket(m_iLocalSockFd);
#else
        ::close(m_iLocalSockFd);
#endif
        m_iLocalSockFd = -1;
    }
    return;
}
